import torch
from torch import nn
from torch.nn import functional as F

class MLP(nn.Module):             # 定义 MLP 类
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(20, 256)   # 定义隐藏层层，输入尺寸为 20，输出尺寸为 256
        self.output = nn.Linear(256, 10)   # 定义输出层，输入尺寸为 256，输出尺寸为 10

    def forward(self, x):          # 定义前向传播函数
        return self.output(F.relu(self.hidden(x)))  # 使用 ReLU 激活函数，计算隐藏层和输出层的输出

net = MLP()

# 存储模型
torch.save(net.state_dict(), 'model.pt')

# 导入模型
model = torch.load('model.pt')
print(model)